# -*- coding: utf-8 -*-
# File : es_util.py
# Author: taoyahui
# Date : 2021/4/1
import os
import sys
sys.path.append('..')
import requests
import json
import urllib3
from utils import common_util
from utils import mysql_util

urllib3.disable_warnings()

json_headers = {
    "Content-Type": "application/json; charset=UTF-8"
}


def get_wazuh_token(ip, username, passwd):
    token = ''
    try:
        # response = requests.get(f'https://47.118.55.203:55000/security/user/authenticate', auth=('wazuh', 'wazuh'), verify=False)
        response = requests.get(f'https://{ip}:55000/security/user/authenticate', auth=(username, passwd),
                                verify=False)
        token = json.loads(response.content)['data']['token']
    except Exception as e:
        print('Get Token failed , reason is ', e)
    return token


def get_distinct_ip(ip, username, passwd):
    try:
        headers = {
            "Content-Type": "application/json; charset=UTF-8"
        }
        payload = {
            'size': 0,
            'aggs': {
                'distinct_ip': {
                    "terms": {
                        "field": "agent.ip"
                    }
                }
            }
        }
        alert_datetime = common_util.current_time_format()
        response = requests.post(f'https://{ip}:9200/wazuh-alerts-4.x-{alert_datetime}/_search', headers=headers,
                                 data=json.dumps(payload), auth=(username, passwd), verify=False)
        res = json.loads(response.content)
        return [data['key'] for data in res['aggregations']['distinct_ip']['buckets']]
    except Exception as e:
        print(e)
    return None


def get_cve_detail(ip, username, passwd, cve, get_ip):
    try:
        payload = {
            # '_source': ['data.vulnerability'],
            'query': {
                'bool': {
                    'must': [{
                        'match': {
                            'data.vulnerability.cve': cve
                        }
                    }, {
                        'match': {
                            "agent.ip": get_ip
                        }
                    }
                    ]

                }
            },
            'size': 1
        }
        # print(payload)
        alert_datetime = common_util.current_time_format()
        response = requests.post(f'https://{ip}:9200/wazuh-alerts-4.x-{alert_datetime}/_search', headers=json_headers,
                                 data=json.dumps(payload), auth=(username, passwd), verify=False)
        res = json.loads(response.content)
        print(res)
        return res
    except Exception as e:
        print(e)


# 通过指定ip获取漏洞信息
def get_cve_alerts(ip, username, passwd, get_ip, time_ago):
    cve_detail_res_arr = []
    try:
        payload = {
            'query': {
                'bool': {
                    'must': [{
                        'term': {
                            "agent.ip": get_ip
                        },
                    }, {
                        'exists': {
                            'field': 'data.vulnerability.cve'
                        }
                    }],
                    'filter': {
                        'range': {
                            '@timestamp': {
                                'gt': f'now-{time_ago}'
                            }
                        }
                    }
                }
            },
            'size': 0,
            'aggs': {
                'distinct_cve': {
                    "terms": {
                        "field": "data.vulnerability.cve",
                        'size': 10000
                    }
                }
            }
        }
        alert_datetime = common_util.current_time_format()
        # alert_datetime = '2021.04.06'
        response = requests.post(f'https://{ip}:9200/wazuh-alerts-4.x-{alert_datetime}/_search', headers=json_headers,
                                 data=json.dumps(payload), auth=(username, passwd), verify=False)
        res = json.loads(response.content)
        print(res)
        print('*' * 20)
        cves = [value['key'] for value in res['aggregations']['distinct_cve']['buckets']]
        print(cves.__len__())
        print('*' * 20)
        print(cves)
        print('*' * 20)
        for cve in cves:
            cve_detail_res_arr.append(get_cve_detail(ip, username, passwd, cve, get_ip))
    except Exception as e:
        print(e)
    return cve_detail_res_arr


def get_alerts(ip, username, passwd):
    payload = {
        'query': {
            'exists': {
                'field': 'data.vulnerability.cve'
            }
        },
        'size': 100
    }
    alert_datetime = common_util.current_time_format()
    response = requests.post(f'https://{ip}:9200/wazuh-alerts-4.x-{alert_datetime}/_search', headers=json_headers,
                             data=json.dumps(payload), auth=(username, passwd), verify=False)
    res = json.loads(response.content)
    print(res)


def get_all_indices(ip, username, passwd):
    response = requests.get(f'https://{ip}:9200/_cat/indices', headers=json_headers, auth=(username, passwd),
                            verify=False)
    # res = json.loads(response.content)
    print(response.content)


def sql_test(ip, username, passwd):
    payload = {
        'query': 'select * from wazuh-alerts-4.x-2021.04.05'
    }
    response = requests.get(f'https://{ip}:9200/_sql?format=txt', headers=json_headers, data=json.dumps(payload),
                            auth=(username, passwd), verify=False)
    print(response.content)


def demo_req():
    get_ip = "192.168.200.132"
    ip = '47.118.55.203'
    username = 'admin'
    passwd = 'admin'
    payload = {
        'query': {
            'bool': {
                'must': [
                    {
                        'term': {
                            f"agent.ip": f"{get_ip}"
                        },
                    },
                    {
                        'term': {

                        }
                    }
                    # {
                    #     'term': {
                    #         'rule.id': 23506
                    #     }
                    # }

                ],
                # 'must': {
                #
                # },
                'filter': {
                    'range': {
                        '@timestamp': {
                            'gt': 'now-1h'
                        }
                    }
                }
            }
        },
        'size': 0,
        'aggs': {
            'distinct_cve': {
                "terms": {
                    "field": "data.vulnerability.cve",
                    'size': 20
                }
            }
        }
    }
    alert_datetime = common_util.current_time_format()
    response = requests.post(f'https://{ip}:9200/wazuh-alerts-4.x-{alert_datetime}/_search', headers=json_headers,
                             data=json.dumps(payload), auth=(username, passwd), verify=False)
    res = json.loads(response.content)
    print(res)
    return res


if __name__ == '__main__':
    # print(get_cve_alerts('47.118.55.203', 'admin', 'admin', '192.168.200.129'))
    # get_distinct_ip('47.118.55.203', 'admin', 'admin')
    # get_alerts('47.118.55.203', 'admin', 'admin')
    get_cve_alerts('47.118.55.203', 'admin', 'SecretPassword', '192.168.200.134')
    # demo_req()

    # get_all_indices('47.118.55.203', 'admin', 'admin')
    # sql_test('47.118.55.203', 'admin', 'admin')
